package zhxt.helmet.utils;
import java.util.*;

public class BoundingBoxTracker {

    private int currentFrameId = 1;
    private int nextTrackId = 1;
    private List<TrackedObject> activeTracks = new ArrayList<>();
    private int maxFrameGap = 3;  // 最大允许的帧间隔
    // 检测框类
    public static class BoundingBox {

        int boxId;         // 框ID（每帧内唯一）
        String label;       // 类别标签
        double x, y;        // 中心点坐标
        double width, height; // 宽高
        double confidence;  // 置信度

        public BoundingBox( int boxId, String label,
                           double x, double y, double width, double height,
                           double confidence) {

            this.boxId = boxId;
            this.label = label;
            this.x = x;
            this.y = y;
            this.width = width;
            this.height = height;
            this.confidence = confidence;
        }
    }

    // 跟踪目标类
    public static class TrackedObject {
        int trackId;                // 跟踪ID
        String label;               // 类别标签
        List<BoundingBox> history;  // 历史检测框
        int lastSeenFrame;          // 最后出现的帧

        public TrackedObject(int trackId, int currentFrameId ,BoundingBox initialBox) {
            this.trackId = trackId;
            this.label = initialBox.label;
            this.history = new ArrayList<>();
            this.history.add(initialBox);
            this.lastSeenFrame = currentFrameId;
        }
        public String getLabel(){
            return label;
        }

        // 更新跟踪目标
        public void update(BoundingBox box, int currentFrameId) {
            this.history.add(box);
            this.lastSeenFrame = currentFrameId;
        }

        // 预测当前位置（简单线性预测）
        public double[] predictPosition() {
            if (history.size() < 2) {
                BoundingBox last = history.get(history.size()-1);
                return new double[]{last.x, last.y};
            }

            BoundingBox last = history.get(history.size()-1);
            BoundingBox prev = history.get(history.size()-2);

            double dx = last.x - prev.x;
            double dy = last.y - prev.y;

            return new double[]{last.x + dx, last.y + dy};
        }
    }


    // 计算两个框的IoU（交并比）
    private double calculateIoU(BoundingBox box1, BoundingBox box2) {
        // 计算相交区域
        double x1 = Math.max(box1.x - box1.width/2, box2.x - box2.width/2);
        double y1 = Math.max(box1.y - box1.height/2, box2.y - box2.height/2);
        double x2 = Math.min(box1.x + box1.width/2, box2.x + box2.width/2);
        double y2 = Math.min(box1.y + box1.height/2, box2.y + box2.height/2);

        double intersection = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);

        // 计算并集
        double area1 = box1.width * box1.height;
        double area2 = box2.width * box2.height;
        double union = area1 + area2 - intersection;

        return union > 0 ? intersection / union : 0;
    }

    // 计算检测框与跟踪目标的相似度
    private double calculateSimilarity(BoundingBox box, TrackedObject track) {
        // 类别不匹配则相似度为0
        if (!box.label.equals(track.label)) {
            return 0;
        }

        // 帧间隔太大则相似度为0
        if (currentFrameId- track.lastSeenFrame > maxFrameGap) {
            return 0;
        }

        // 计算IoU
        BoundingBox lastBox = track.history.get(track.history.size()-1);
        double iou = calculateIoU(box, lastBox);

        // 简单运动模型预测
        double[] predictedPos = track.predictPosition();
        double dist = Math.sqrt(Math.pow(box.x - predictedPos[0], 2) +
                Math.pow(box.y - predictedPos[1], 2));
        double distScore = 1.0 / (1.0 + dist);

        // 综合相似度
        return 0.7 * iou + 0.3 * distScore;
    }

    // 使用匈牙利算法进行匹配
    private int[] hungarianMatch(double[][] costMatrix) {
        int rows = costMatrix.length;
        int cols = rows > 0 ? costMatrix[0].length : 0;

        // 这里简化为贪婪匹配，实际应用中应实现完整匈牙利算法
        int[] assignment = new int[rows];
        Arrays.fill(assignment, -1);

        Set<Integer> assignedCols = new HashSet<>();

        for (int i = 0; i < rows; i++) {
            int bestJ = -1;
            double bestScore = -1;

            for (int j = 0; j < cols; j++) {
                if (!assignedCols.contains(j) && costMatrix[i][j] > bestScore) {
                    bestScore = costMatrix[i][j];
                    bestJ = j;
                }
            }

            if (bestJ != -1 && bestScore > 0) {
                assignment[i] = bestJ;
                assignedCols.add(bestJ);
            }
        }

        return assignment;
    }

    // 处理一帧的检测结果, 并返回新目标
    public List<TrackedObject> processFrame(List<BoundingBox> frameBoxes) {
        List<TrackedObject> newTracks = new ArrayList<>();
        if (frameBoxes.isEmpty()) {
            // 没有检测框，清理过期的跟踪
            activeTracks.removeIf(track ->
                    currentFrameId - track.lastSeenFrame > maxFrameGap);
            currentFrameId++;
            return new ArrayList<>(newTracks);
        }


        // 1. 构建相似度矩阵
        double[][] similarityMatrix = new double[frameBoxes.size()][activeTracks.size()];
        for (int i = 0; i < frameBoxes.size(); i++) {
            for (int j = 0; j < activeTracks.size(); j++) {
                similarityMatrix[i][j] = calculateSimilarity(frameBoxes.get(i), activeTracks.get(j));
            }
        }

        // 2. 进行匹配
        int[] assignment = hungarianMatch(similarityMatrix);

        Set<Integer> matchedTrackIndices = new HashSet<>();
        Set<Integer> matchedBoxIndices = new HashSet<>();

        // 3. 更新匹配的跟踪
        for (int i = 0; i < assignment.length; i++) {
            if (assignment[i] != -1) {
                activeTracks.get(assignment[i]).update(frameBoxes.get(i), currentFrameId);
                matchedTrackIndices.add(assignment[i]);
                matchedBoxIndices.add(i);
            }
        }

        // 4. 为未匹配的检测创建新跟踪
        for (int i = 0; i < frameBoxes.size(); i++) {
            if (!matchedBoxIndices.contains(i)) {
                TrackedObject newTrack = new TrackedObject(nextTrackId++, currentFrameId, frameBoxes.get(i));
                newTracks.add(newTrack);
                activeTracks.add(newTrack);
            }
        }

        // 5. 清理未匹配且过期的跟踪
        List<TrackedObject> toRemove = new ArrayList<>();
        for (int j = 0; j < activeTracks.size(); j++) {
            if (!matchedTrackIndices.contains(j) &&
                    currentFrameId - activeTracks.get(j).lastSeenFrame > maxFrameGap) {
                toRemove.add(activeTracks.get(j));
            }
        }
        activeTracks.removeAll(toRemove);

        currentFrameId++;

        return new ArrayList<>(newTracks);
    }

    public static void main(String[] args) {
        // 示例用法
        BoundingBoxTracker tracker = new BoundingBoxTracker();

        // 模拟一些帧数据
        List<List<BoundingBox>> frames = new ArrayList<>();

        // 第1帧
        List<BoundingBox> frame1 = new ArrayList<>();
        frame1.add(new BoundingBox( 1, "person", 10, 10, 5, 10, 0.9));
        frame1.add(new BoundingBox( 2, "car", 50, 50, 20, 15, 0.8));
        frames.add(frame1);

        // 第2帧
        List<BoundingBox> frame2 = new ArrayList<>();
        frame2.add(new BoundingBox(  1, "person", 12, 11, 5, 10, 0.85));
        frame2.add(new BoundingBox(  2, "car", 52, 52, 20, 15, 0.82));
        frame2.add(new BoundingBox(  3, "dog", 30, 30, 8, 8, 0.7));
        frames.add(frame2);

        // 第3帧
        List<BoundingBox> frame3 = new ArrayList<>();
        frame3.add(new BoundingBox(  1, "person", 14, 12, 5, 10, 0.88));
        frame3.add(new BoundingBox(  2, "car", 54, 54, 20, 15, 0.85));
        frames.add(frame3);
// 第4帧
        List<BoundingBox> frame4 = new ArrayList<>();
        frame4.add(new BoundingBox(  1, "person", 12, 12, 5, 10, 0.88));
        frame4.add(new BoundingBox(  2, "car", 56, 54, 20, 15, 0.85));
        frames.add(frame4);
        //空帧
        List<BoundingBox> ft = new ArrayList<>();
        frames.add(ft);
        frames.add(ft);
        frames.add(ft);
        // 第5帧
        List<BoundingBox> frame5 = new ArrayList<>();
        frame5.add(new BoundingBox(  1, "person", 12, 12, 5, 10, 0.88));
        frame5.add(new BoundingBox(  2, "car", 56, 54, 20, 15, 0.85));
        frames.add(frame5);
        // 第6 帧
        List<BoundingBox> frame6 = new ArrayList<>();
        frame6.add(new BoundingBox(  1, "person", 12, 12, 5, 10, 0.88));
        frame6.add(new BoundingBox(  2, "car", 56, 54, 20, 15, 0.85));
        frames.add(frame6);
        // 第7帧
        List<BoundingBox> frame7 = new ArrayList<>();
        frame7.add(new BoundingBox( 1, "a", 12, 12, 5, 10, 0.88));
        frame7.add(new BoundingBox(  2, "car", 56, 54, 20, 15, 0.85));
        frames.add(frame7);
        // 处理所有帧
        for (List<BoundingBox> frame : frames) {
            List<TrackedObject> tracks = tracker.processFrame(frame);
            System.out.println("Frame " + (tracker.currentFrameId-1) + " tracks:");
            for (TrackedObject track : tracks) {
                BoundingBox lastBox = track.history.get(track.history.size()-1);
                System.out.printf("Track %d (%s): (%.1f, %.1f)%n",
                        track.trackId, track.label, lastBox.x, lastBox.y);
            }
            System.out.println();
        }
    }
}